#include "BartClassifier.h"

using namespace bart;

std::string BartClassifier::model_path = "";
std::shared_ptr<BartClassifier> BartClassifier::classifier = nullptr;

BartClassifier::BartClassifier(const std::string &model_path) {
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    bart_classifier = torch::jit::load(model_path);
  } catch (const c10::Error &e) {
    std::cerr << "error loading the model\n";
  }
}

std::shared_ptr<BartClassifier>
BartClassifier::getInstance(const std::string &model_path) {
  if (model_path != BartClassifier::model_path) {
    BartClassifier::classifier = std::make_shared<BartClassifier>(model_path);
    // add the logic of clean classifier
    BartClassifier::model_path = model_path;
  }
  return BartClassifier::classifier;
}

BartClassifier::~BartClassifier() {}

at::Tensor BartClassifier::getPossibility(const std::vector<int> &encoding) {
  std::vector<torch::jit::IValue> inputs;
  auto options = torch::TensorOptions().dtype(torch::kLong);
  auto input = torch::full({1, long(encoding.size())}, 0, options);
  for (int i = 0; i < encoding.size(); ++i) {
    input[0][i] = encoding.at(i);
  }
  inputs.push_back(input);
  auto logits = this->bart_classifier.forward(inputs);
  if (logits.isTuple()) {
    auto logits_tuple = logits.toTuple();
    auto logits_first = logits_tuple->elements().at(0).toTensor();
    auto entail_contradiction_logits = logits_first.slice(1, 0, 3, 2);
    auto probs = entail_contradiction_logits.softmax(1);
    auto prob_is_true = probs.slice(1, 1, 2, 1);
    std::cout << prob_is_true << std::endl;
    return prob_is_true;
  } else {
    exit(-1);
  }
}